import asyncio
import json
import ssl
from websockets.server import serve
from websockets.exceptions import ConnectionClosed

# 存储连接的客户端
clients = {}

async def handle_client(websocket, path):
    client_id = None
    try:
        # 接收客户端ID
        client_id = await websocket.recv()
        clients[client_id] = websocket
        print(f"客户端 {client_id} 已连接")
        
        # 发送在线客户端列表
        peer_list = [pid for pid in clients.keys() if pid != client_id]
        await websocket.send(json.dumps({
            "type": "peers",
            "peers": peer_list
        }))

        # 接收并转发消息
        async for message in websocket:
            try:
                data = json.loads(message)
                
                # 处理不同类型的消息
                if data.get("target"):
                    if data["target"] in clients:
                        # 转发消息给指定客户端
                        await clients[data["target"]].send(message)
                        print(f"消息从 {client_id} 转发到 {data['target']}")
                    else:
                        # 目标客户端不在线，发送错误消息
                        await websocket.send(json.dumps({
                            "type": "error",
                            "message": f"目标客户端 {data['target']} 不在线"
                        }))
            except json.JSONDecodeError as e:
                print(f"JSON解析错误: {e}")
                await websocket.send(json.dumps({
                    "type": "error",
                    "message": "消息格式错误"
                }))
            except Exception as e:
                print(f"处理消息时出错: {e}")
                
    except ConnectionClosed:
        print(f"客户端 {client_id} 连接已关闭")
    except Exception as e:
        print(f"处理客户端 {client_id} 时出错: {e}")
    finally:
        if client_id and client_id in clients:
            del clients[client_id]
            # 通知其他客户端
            disconnected_clients = []
            for other_id, other_ws in clients.items():
                try:
                    await other_ws.send(json.dumps({
                        "type": "peer_disconnected",
                        "peer": client_id
                    }))
                except Exception as e:
                    print(f"通知客户端 {other_id} 时出错: {e}")
                    disconnected_clients.append(other_id)
            
            # 清理断开的客户端
            for disconnected_id in disconnected_clients:
                if disconnected_id in clients:
                    del clients[disconnected_id]
                    
        print(f"客户端 {client_id} 已断开连接")

async def main():
    ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    ssl_context.load_cert_chain('certificate/server.crt', 'certificate/server.key')
    
    async with serve(handle_client, "0.0.0.0", 8765, ssl=ssl_context):
        print("WebSocket信令服务器已启动，监听端口 8765")
        await asyncio.Future()  # 保持服务器运行

if __name__ == "__main__":
    asyncio.run(main())